import cv2
import numpy as np
import logging
import os
import sys

# 导入新的封装类
from PaddleOcrSealDetector import PaddleOcrSealDetector
from PaddleXSealDetector import PaddleXSealDetector

# 配置日志
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

class SealRecognizer:
    def __init__(self, config_path=None, use_paddlex=False, use_paddleocr=True, use_paddleocr_handwritten=False, pdx_det_model_dir=None, pdx_rec_model_dir=None, use_handwritten=False):
        """初始化：加载PaddleOCR模型，可选加载PaddleX模型和手写体识别"""
        
        # 初始化PaddleOCR检测器（可选）
        self.use_paddleocr = use_paddleocr
        self.use_paddleocr_handwritten = use_paddleocr_handwritten
        self.paddleocr_detector = None
        
        if use_paddleocr:
            try:
                self.paddleocr_detector = PaddleOcrSealDetector(use_handwritten=use_paddleocr_handwritten)
                logger.info("PaddleOCR检测器初始化成功")
                if use_paddleocr_handwritten:
                    logger.info("PaddleOCR手写体识别功能已启用")
            except Exception as e:
                logger.error(f"PaddleOCR检测器初始化失败：{str(e)}", exc_info=True)
                logger.warning("PaddleOCR初始化失败，将仅使用PaddleX进行识别")
                self.use_paddleocr = False
        else:
            logger.info("PaddleOCR识别已禁用")
        
        # 初始化PaddleX检测器（可选）
        self.use_paddlex = use_paddlex
        self.paddlex_detector = None
        
        if use_paddlex:
            try:
                self.paddlex_detector = PaddleXSealDetector(config_path, use_handwritten)
                if self.paddlex_detector.pdx_pipeline is None:
                    logger.warning("PaddleX模型未成功加载，将仅使用PaddleOCR")
                    self.use_paddlex = False
                else:
                    logger.info("PaddleX检测器初始化成功")
            except Exception as e:
                logger.warning(f"PaddleX检测器初始化失败：{str(e)}，将仅使用PaddleOCR")
                self.use_paddlex = False

    def detect_seal_region(self, img_path):
        """
        印章区域检测：结合传统OpenCV方法和PaddleX方法，取并集
        返回：(图章区域列表, 原图, 红色掩码图)
        """
        # 1. 读取图片并校验
        img = cv2.imread(img_path)
        if img is None:
            error_msg = f"无法读取图片：路径错误或文件损坏（路径：{img_path}）"
            logger.error(error_msg)
            raise ValueError(error_msg)
        logger.info(f"成功读取图片，尺寸：{img.shape[0]}x{img.shape[1]}")

        # 2. 传统OpenCV方法检测（如果启用）
        opencv_regions = []
        if self.use_paddleocr and self.paddleocr_detector:
            opencv_regions, _, _ = self.paddleocr_detector.detect_seal_region(img_path)
        
        # 3. PaddleX方法检测（如果启用）
        paddlex_regions = []
        if self.use_paddlex and self.paddlex_detector:
            paddlex_regions, _, _ = self.paddlex_detector.detect_seal_region(img_path)
        
        # 4. 合并所有检测结果（去重）
        all_regions = []
        if opencv_regions:
            all_regions.extend(opencv_regions)
        if paddlex_regions:
            all_regions.extend(paddlex_regions)
        
        # 去重：基于区域重叠度
        unique_regions = self._merge_overlapping_regions(all_regions)
        
        # 如果启用了PaddleOCR手写体识别，检测手写体区域并标注到seal_roi_marked.jpg
        if self.use_paddleocr_handwritten:
            self._add_handwritten_regions_to_marked_image(img_path, img)
        
        logger.info(f"检测到印章区域总数：{len(unique_regions)}（OpenCV: {len(opencv_regions) if opencv_regions else 0}, PaddleX: {len(paddlex_regions) if paddlex_regions else 0}）")
        
        return unique_regions, img, None

    def _merge_overlapping_regions(self, regions):
        """合并重叠的检测区域"""
        if not regions:
            return []
        
        # 计算两个区域的重叠度
        def overlap_ratio(region1, region2):
            x1, y1, w1, h1 = region1
            x2, y2, w2, h2 = region2
            
            # 计算交集
            x_left = max(x1, x2)
            y_top = max(y1, y2)
            x_right = min(x1 + w1, x2 + w2)
            y_bottom = min(y1 + h1, y2 + h2)
            
            if x_right < x_left or y_bottom < y_top:
                return 0.0
            
            intersection = (x_right - x_left) * (y_bottom - y_top)
            union = w1 * h1 + w2 * h2 - intersection
            return intersection / union if union > 0 else 0.0
        
        # 合并重叠区域
        merged_regions = []
        used = [False] * len(regions)
        
        for i, region1 in enumerate(regions):
            if used[i]:
                continue
                
            # 找到所有与当前区域重叠的区域
            overlapping = [region1]
            for j, region2 in enumerate(regions[i+1:], i+1):
                if used[j]:
                    continue
                if overlap_ratio(region1, region2) > 0.3:  # 重叠度阈值
                    overlapping.append(region2)
                    used[j] = True
            
            # 计算合并后的区域（取最大外接矩形）
            if len(overlapping) > 1:
                all_x = [r[0] for r in overlapping]
                all_y = [r[1] for r in overlapping]
                all_w = [r[2] for r in overlapping]
                all_h = [r[3] for r in overlapping]
                
                x_min = min(all_x)
                y_min = min(all_y)
                x_max = max(x + w for x, w in zip(all_x, all_w))
                y_max = max(y + h for y, h in zip(all_y, all_h))
                
                merged_region = (x_min, y_min, x_max - x_min, y_max - y_min)
                merged_regions.append(merged_region)
                logger.info(f"合并了{len(overlapping)}个重叠区域：{merged_region}")
            else:
                merged_regions.append(region1)
            
            used[i] = True
        
        return merged_regions

    def _add_handwritten_regions_to_marked_image(self, img_path, img):
        """将PaddleOCR手写体区域添加到seal_roi_marked.jpg中"""
        try:
            # 如果PaddleOCR检测器不存在，创建一个专门用于手写体识别的检测器
            if not self.paddleocr_detector:
                from PaddleOcrSealDetector import PaddleOcrSealDetector
                self.paddleocr_detector = PaddleOcrSealDetector(use_handwritten=True)
                logger.info("创建独立的PaddleOCR手写体检测器")
            
            # 检测手写体区域
            handwritten_regions = self.paddleocr_detector.detect_handwritten_regions(img_path)
            
            if handwritten_regions:
                # 尝试读取现有的seal_roi_marked.jpg，如果不存在则使用原图
                marked_path = os.path.join(os.path.dirname(img_path), "seal_roi_marked.jpg")
                if os.path.exists(marked_path):
                    img_marked = cv2.imread(marked_path)
                    logger.info("读取现有的seal_roi_marked.jpg，添加PaddleOCR手写体标注")
                else:
                    img_marked = img.copy()
                    logger.info("创建新的标注图像，添加PaddleOCR手写体标注")
                
                # 使用绿色框标注手写体区域
                for i, region in enumerate(handwritten_regions):
                    x, y, w, h = region['x'], region['y'], region['w'], region['h']
                    # 绘制绿色边界框
                    cv2.rectangle(img_marked, (x, y), (x + w, y + h), (0, 255, 0), 2)  # 绿色框
                    # 添加标签
                    label = f"PaddleOCR-hand-{i+1}"
                    cv2.putText(img_marked, label, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
                
                # 保存合并后的标注图像
                cv2.imwrite(marked_path, img_marked)
                logger.info(f"PaddleOCR手写体区域已添加到：{marked_path}（共{len(handwritten_regions)}个区域）")
            else:
                logger.info("PaddleOCR未检测到手写体区域")
                
        except Exception as e:
            logger.warning(f"添加PaddleOCR手写体区域到标注图像失败：{str(e)}")

    def recognize_seal_text(self, img_path, img, seal_regions):
        """
        识别多个印章区域的文字，分别使用传统方法和PaddleX方法
        返回：识别结果字典
        """
        if not seal_regions:
            logger.warning("未检测到图章区域，无法识别文字")
            return {"opencv_results": [], "paddlex_results": []}

        logger.info(f"开始识别{len(seal_regions)}个图章区域的文字")
        
        # 1. 传统OpenCV方法识别（如果启用）
        opencv_results = []
        if self.use_paddleocr and self.paddleocr_detector:
            opencv_results = self.paddleocr_detector.recognize_seal_text(img_path, img, seal_regions)
        
        # 2. PaddleX方法识别（如果启用）
        paddlex_results = []
        if self.use_paddlex and self.paddlex_detector:
            paddlex_results = self.paddlex_detector.recognize_seal_text(img_path, img, seal_regions)

        return {
            "opencv_results": opencv_results,
            "paddlex_results": paddlex_results
        }

    def recognize_full_text(self, img_path):
        """识别图片中的全部文字"""
        if self.use_paddleocr and self.paddleocr_detector:
            return self.paddleocr_detector.recognize_full_text(img_path)
        else:
            logger.warning("PaddleOCR未启用，无法进行全文识别")
            return "PaddleOCR未启用，无法进行全文识别"

    def detect_handwritten_text(self, img_path, seal_regions=None):
        """检测手写体文字区域"""
        if not self.use_paddlex or not self.paddlex_detector:
            logger.info("手写体识别未启用或PaddleX不可用")
            return []
        
        return self.paddlex_detector.detect_handwritten_text(img_path, seal_regions)

    def recognize_handwritten_text(self, img, handwritten_regions):
        """识别手写体文字"""
        if not self.use_paddlex or not self.paddlex_detector:
            return handwritten_regions
        
        return self.paddlex_detector.recognize_handwritten_text(img, handwritten_regions)

    def visualize_handwritten_regions(self, img, handwritten_regions, img_path):
        """可视化手写体区域"""
        if not self.use_paddlex or not self.paddlex_detector:
            return None
        
        return self.paddlex_detector.visualize_handwritten_regions(img, handwritten_regions, img_path)

    def detect_paddleocr_handwritten_regions(self, img_path):
        """使用PaddleOCR检测手写体区域"""
        if not self.use_paddleocr_handwritten:
            logger.info("PaddleOCR手写体识别未启用")
            return []
        
        # 如果PaddleOCR检测器不存在，创建一个专门用于手写体识别的检测器
        if not self.paddleocr_detector:
            from PaddleOcrSealDetector import PaddleOcrSealDetector
            self.paddleocr_detector = PaddleOcrSealDetector(use_handwritten=True)
            logger.info("创建独立的PaddleOCR手写体检测器")
        
        return self.paddleocr_detector.detect_handwritten_regions(img_path)

    def recognize_paddleocr_handwritten_text(self, img, handwritten_regions):
        """使用PaddleOCR识别手写体文字"""
        if not self.use_paddleocr_handwritten:
            return handwritten_regions
        
        # 如果PaddleOCR检测器不存在，创建一个专门用于手写体识别的检测器
        if not self.paddleocr_detector:
            from PaddleOcrSealDetector import PaddleOcrSealDetector
            self.paddleocr_detector = PaddleOcrSealDetector(use_handwritten=True)
            logger.info("创建独立的PaddleOCR手写体检测器")
        
        return self.paddleocr_detector.recognize_handwritten_text(img, handwritten_regions)

    def visualize_paddleocr_handwritten_regions(self, img, handwritten_regions, img_path):
        """可视化PaddleOCR手写体区域"""
        if not self.use_paddleocr_handwritten:
            return None
        
        # 如果PaddleOCR检测器不存在，创建一个专门用于手写体识别的检测器
        if not self.paddleocr_detector:
            from PaddleOcrSealDetector import PaddleOcrSealDetector
            self.paddleocr_detector = PaddleOcrSealDetector(use_handwritten=True)
            logger.info("创建独立的PaddleOCR手写体检测器")
        
        return self.paddleocr_detector.visualize_handwritten_regions(img, handwritten_regions, img_path)

    def recognize(self, img_path):
        """完整识别流程：定位→识别→结果返回"""
        logger.info("=" * 50)
        logger.info(f"开始图章识别：图片路径={img_path}")
        try:
            # 保存当前图片路径供PaddleX使用
            self._current_img_path = img_path
            
            # 1. 定位图章区域（返回多个区域）
            seal_regions, img, red_mask = self.detect_seal_region(img_path)
            
            # 2. 识别全部内容
            full_text = self.recognize_full_text(img_path)
            
            # 3. 手写体识别（如果启用）
            handwritten_info = ""
            paddleocr_handwritten_info = ""
            
            # PaddleX手写体识别
            if self.use_paddlex and self.paddlex_detector:
                handwritten_regions = self.detect_handwritten_text(img_path, seal_regions)
                if handwritten_regions:
                    # 识别手写体文字
                    handwritten_regions = self.recognize_handwritten_text(img, handwritten_regions)
                    
                    # 可视化手写体区域
                    self.visualize_handwritten_regions(img, handwritten_regions, img_path)
                    
                    # 格式化手写体信息
                    handwritten_texts = []
                    for region in handwritten_regions:
                        x, y, w, h = region['x'], region['y'], region['w'], region['h']
                        text = region.get('text', '')
                        detection_conf = region['detection_confidence']
                        recognition_conf = region.get('recognition_confidence', 0.0)
                        handwritten_texts.append(f"PaddleX手写体区域{region['region_id']}: 位置({x},{y}), 尺寸({w}x{h}), 检测置信度{detection_conf:.3f}, 识别置信度{recognition_conf:.3f}, 内容: {text}")
                    
                    handwritten_info = f"PaddleX手写体识别结果：\n" + "\n".join(handwritten_texts)
                else:
                    handwritten_info = "PaddleX未检测到手写体文字"
            
            # PaddleOCR手写体识别（独立工作）
            if self.use_paddleocr_handwritten:
                paddleocr_handwritten_regions = self.detect_paddleocr_handwritten_regions(img_path)
                if paddleocr_handwritten_regions:
                    # 识别手写体文字
                    paddleocr_handwritten_regions = self.recognize_paddleocr_handwritten_text(img, paddleocr_handwritten_regions)
                    
                    # 可视化手写体区域
                    self.visualize_paddleocr_handwritten_regions(img, paddleocr_handwritten_regions, img_path)
                    
                    # 格式化手写体信息
                    paddleocr_handwritten_texts = []
                    for region in paddleocr_handwritten_regions:
                        x, y, w, h = region['x'], region['y'], region['w'], region['h']
                        text = region.get('text', '')
                        detection_conf = region['detection_confidence']
                        recognition_conf = region.get('recognition_confidence', 0.0)
                        paddleocr_handwritten_texts.append(f"PaddleOCR手写体区域{region['region_id']}: 位置({x},{y}), 尺寸({w}x{h}), 检测置信度{detection_conf:.3f}, 识别置信度{recognition_conf:.3f}, 内容: {text}")
                    
                    paddleocr_handwritten_info = f"PaddleOCR手写体识别结果：\n" + "\n".join(paddleocr_handwritten_texts)
                else:
                    paddleocr_handwritten_info = "PaddleOCR未检测到手写体文字"
            
            # 4. 处理图章信息
            if not seal_regions:
                seal_info = "未检测到图章区域（可查看red_mask_debug.jpg调试）"
                paddlex_info = ""
            else:
                # 使用新的识别方法，分别获取OpenCV和PaddleX结果
                recognition_results = self.recognize_seal_text(img_path, img, seal_regions)
                
                # 处理OpenCV结果
                opencv_results = recognition_results.get("opencv_results", [])
                if opencv_results:
                    opencv_texts = []
                    for result in opencv_results:
                        text = result["text"]
                        opencv_texts.append(f"OpenCV识别: {text}")
                    seal_info = f"OpenCV识别结果：\n" + "\n".join(opencv_texts)
                else:
                    seal_info = "OpenCV未识别到图章文字"
                
                # 处理PaddleX结果
                paddlex_results = recognition_results.get("paddlex_results", [])
                if paddlex_results:
                    paddlex_texts = []
                    for result in paddlex_results:
                        region = result["region"]
                        text = result["text"]
                        x, y, w, h = region
                        paddlex_texts.append(f"区域({x},{y},{w},{h}): {text}")
                    paddlex_info = f"PaddleX识别结果：\n" + "\n".join(paddlex_texts)
                else:
                    if self.use_paddlex:
                        paddlex_info = "PaddleX未识别到图章文字"
                    else:
                        paddlex_info = "PaddleX未启用"

            # 5. 整合结果
            result_parts = [f"全部内容识别：{full_text}"]
            result_parts.append(f"\n{seal_info}")
            if self.use_paddlex:
                result_parts.append(f"\n{paddlex_info}")
            if self.use_paddlex and handwritten_info:
                result_parts.append(f"\n{handwritten_info}")
            if self.use_paddleocr_handwritten and paddleocr_handwritten_info:
                result_parts.append(f"\n{paddleocr_handwritten_info}")
            
            return "\n".join(result_parts)

        except Exception as e:
            error_msg = f"识别过程异常：{str(e)}"
            logger.error(error_msg, exc_info=True)
            return error_msg


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("使用方法: python seal_recognizer.py <图片路径> [--paddleocr] [--paddleocrhand] [--paddlex] [--handwritten]")
        print("选项:")
        print("  --paddleocr            启用PaddleOCR识别")
        print("  --paddleocrhand        启用PaddleOCR手写体识别")
        print("  --paddlex              启用PaddleX识别对比（使用Hugging Face模型）")
        print("  --handwritten          启用手写体识别（需要PaddleX）")
        print("")
        print("示例:")
        print("  # 仅PaddleOCR（默认）")
        print("  python seal_recognizer.py image.jpg")
        print("")
        print("  # 仅PaddleOCR手写体识别（独立使用）")
        print("  python seal_recognizer.py image.jpg --paddleocrhand")
        print("")
        print("  # PaddleOCR + 手写体识别")
        print("  python seal_recognizer.py image.jpg --paddleocr --paddleocrhand")
        print("")
        print("  # 仅PaddleX，不使用PaddleOCR")
        print("  python seal_recognizer.py image.jpg --paddlex")
        print("")
        print("  # 同时使用PaddleOCR和PaddleX")
        print("  python seal_recognizer.py image.jpg --paddleocr --paddlex")
        print("")
        print("  # 仅PaddleX + 手写体识别")
        print("  python seal_recognizer.py image.jpg --paddlex --handwritten")
        print("")
        print("  # 全部功能")
        print("  python seal_recognizer.py image.jpg --paddleocr --paddleocrhand --paddlex --handwritten")
        print("")
        print("注意：PaddleX使用Hugging Face模型，首次运行会自动下载")
        sys.exit(1)
    
    img_path = sys.argv[1]
    
    # 解析命令行参数
    use_paddleocr = "--paddleocr" in sys.argv
    use_paddleocr_handwritten = "--paddleocrhand" in sys.argv
    use_paddlex = "--paddlex" in sys.argv
    use_handwritten = "--handwritten" in sys.argv
    
    # 如果没有指定任何识别方法，默认使用PaddleOCR
    if not use_paddleocr and not use_paddlex and not use_paddleocr_handwritten:
        use_paddleocr = True
        print("未指定识别方法，默认启用PaddleOCR")
    elif use_paddleocr_handwritten and not use_paddleocr and not use_paddlex:
        print("仅启用手写体识别，不进行印章检测")
    
    # 如果启用手写体识别，需要同时启用PaddleX
    if use_handwritten and not use_paddlex:
        print("警告：手写体识别需要PaddleX支持，自动启用PaddleX")
        use_paddlex = True
    
    # 如果启用PaddleX，显示模型信息
    if use_paddlex:
        print("=" * 60)
        print("PaddleX Hugging Face模型配置")
        print("=" * 60)
        print("检测模型: PaddlePaddle/PaddleX-Seal-Detection")
        print("识别模型: PaddlePaddle/PaddleX-Seal-Recognition")
        if use_handwritten:
            print("手写体识别: 已启用")
        print("首次运行将自动下载模型")
        print("=" * 60)
    
    # 创建识别器
    recognizer = SealRecognizer(
        use_paddlex=use_paddlex, 
        use_paddleocr=use_paddleocr, 
        use_paddleocr_handwritten=use_paddleocr_handwritten,
        use_handwritten=use_handwritten
    )
    
    result = recognizer.recognize(img_path)
    print("\n" + "=" * 50)
    print(result)
    print("=" * 50)